/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010, 2011 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

// sysGen.c -- functions for generating a system (coordinates, velocities, etc.) from scratch.

#include "dataStructs.h"
#include "definitions.h"

extern RMat *siteRotMat;
extern VecR *siteCoords, *siteOrientVec, *siteVelocity, *siteAngMom; 
extern const int nSites, stepCount;
extern const LipidStruct dopcStruct, dopeStruct;
extern Site *site;  
extern const real *inertia, *mass;
extern const VecI initUcell;

/* EXTERNAL FUNCTION DECLARATIONS: */
void BuildRotMatrixFromQuat( RMat*, Quat*, int );
void CheckNLipids( int, int );
void CheckWatInputParameters( int );
void EulerToQuat( Quat*, real* );
void EulerToRotMat( RMat*, real, real, real );
void GetOrientPrincAxisZ_FromRotMat( VecR*, RMat ); 
void PrintMatrix( RMat );
real RandR();
void VRand( VecR* );

// global variables
const VecI initHalfCellWat; // cell array employed to hydrate a lipid bilayer 
static const real INITIAL_BILAYER_THICKNESS = 3.5; /* [nm] initial bilayer thickness set to typical value */

void InitSymmRigBodOrientations( FILE *fPtr ) // Adapted from [Rapaport, pp. 214-215]
{
  int n;
  VecR e;
  real eAng[ 3 ];

  fprintf( fPtr, "InitSymmRigBodOrientations... " ); fflush(fPtr);

  DO_SITE {
    if ( SRB == site[ n ].mechType ) {
      VRand( &e );                   // set orientations at random
      eAng[ 0 ] = atan2( e.x, e.y ); // phi
      eAng[ 1 ] = acos( e.z );       // theta
      eAng[ 2 ] = 2. * PI * RandR(); // psi

      EulerToRotMat( &siteRotMat[ n ], eAng[ 0 ], eAng[ 1 ], eAng[ 2 ] ); 

      GetOrientPrincAxisZ_FromRotMat( &siteOrientVec[ n ], siteRotMat[ n ] ); /* convert rotation matrix to orientation vector */

      //      printf("### %f %f %f %f\n",VLen(siteOrientVec[ n ]),siteOrientVec[ n ].x,siteOrientVec[ n ].y,siteOrientVec[ n ].z);
      /*       VSet( siteOrientVec[ n ], 0, 0, 1 ); /\* Set orientations of SRB sites aligned along the z direction -
	       DISCRIMINATE BETWEEN TOP AND BOTTOM MONOLAYER!*\/ */
      /*       eAng[ 1 ] = atan2( siteOrientVec[ n ].x, siteOrientVec[ n ].y ); */
      /*       eAng[ 2 ] = acos( siteOrientVec[ n ].z ); */
      /*       eAng[ 0 ] = 2. * PI * RandR(); */
      /*       EulerToQuat( &qe, eAng );	     */
    
    }
  }
  fprintf( fPtr, "done\n" ); fflush(fPtr);
}

/* InitCoords -- adapted from Rapaport [p.67,68] distributes the nSites on a Face-Centred Cubic lattice,
   with the option of unequal edge lengths, so that each unit cell contains four
   atoms, and the system is centered about the origin. */
void InitCoords( const VecR region, FILE *fPtr )
{
  VecR c, gap;
  int j, n, nx, ny, nz;

  fprintf( fPtr, "InitCoords... " ); fflush(fPtr);
  
  VDiv( gap, region, initUcell );
  n = 0;
  for ( nz = 0; nz < initUcell.z; nz++ ) {
    for ( ny = 0; ny < initUcell.y; ny++ ) {
      for ( nx = 0; nx < initUcell.x; nx++ ) {
        VSet( c, nx + 0.25, ny + 0.25, nz + 0.25 );
        VMul( c, c, gap );
        VVSAdd( c, -0.5, region );
	for ( j = 0; j < 4; j++ ) {
	  siteCoords[ n ] = c; 
	  if ( j != 3 ) {
	    if ( j != 0 ) siteCoords[ n ].x += 0.5 * gap.x; 
	    if ( j != 1 ) siteCoords[ n ].y += 0.5 * gap.y; 
	    if ( j != 2 ) siteCoords[ n ].z += 0.5 * gap.z; 
	  }
	  ++n;
	}
      }
    }
  }
  fprintf( fPtr, "done\n" ); fflush(fPtr);
}

/* InitVels. Setting of initial COM velocities. */
void InitVels( const real extTemperature, FILE *fPtr )
{
  int n;
  real velSum = 0.;
  //  real velMag = sqrt( 3 * ( 1. - 1. / nSites ) * extTemperature ); 
  real velMag = sqrt( 3 * kB_IN_BRAHMS_UNITS * extTemperature / mass[ WATER ] ); 
  VecR vSum;

  fprintf( fPtr, "InitVels... "); fflush( fPtr );

  VZero( vSum );

  DO_SITE {
    VRand( &siteVelocity[ n ] ); /* randomly oriented unit vector */
    VScale( siteVelocity[ n ], velMag );
    VVAdd( vSum, siteVelocity[ n ] );
  }
  DO_SITE { 
    VVSAdd( siteVelocity[ n ], - 1.0 / nSites, vSum );
    //    printf("***site[ %4d ].v = %f\n", n, VLen( siteVelocity[ n ] ) ); fflush(stdout);
    velSum += VLenSq( siteVelocity[ n ] ); /* just checking */
  }
  fprintf( fPtr, "done\n" ); fflush(fPtr);
}

void InitAngVelsSRB( const real extTemperature, FILE *fPtr ) 
{ /* initialise *linear* angular velocities - from allenTildesley F.24 */
  real f;
  VecR e,w,angMomentum;
  real mean, u_dot_e, uu, uFac;
  int n;

  fprintf( fPtr, "InitAngVelsSRB... " ); fflush(fPtr);

  DO_SITE {
    if ( site[ n ].mechType == SRB ) {
      VRand( &siteAngMom[ n ] );
      u_dot_e = VDot( siteAngMom[ n ], siteOrientVec[ n ] );
      VVSAdd( siteAngMom[ n ], - u_dot_e, siteOrientVec[ n ] );
      uu = VLenSq( siteAngMom[ n ] );
      VScale( siteAngMom[ n ], 1. / ( sqrt( uu ) ) );
      mean = 2. * kB_IN_BRAHMS_UNITS * extTemperature / inertia[ site[ n ].type ];
      uFac = sqrt( - mean * log( RandR() ) );
      VScale( siteAngMom[ n ], uFac );      

      f = sqrt( 2 * kB_IN_BRAHMS_UNITS * extTemperature / ( 3. * inertia[ site[ n ].type ] ) );
      VCross( angMomentum, siteCoords[ n ], siteVelocity[ n ] );
      VRand( &e );
      /*       f = 0.5 * velMag / sqrt (VWLenSq (ssdInertia, e)); */
      VSCopy( w, f, e );                         /* body-fixed angular velocity */ 
      VSCopy( angMomentum, inertia[ site[ n ].type ], w );        /* body-fixed angular momentum */
      VCopy( siteAngMom[ n ], angMomentum ); /* body-fixed angular momenta */


      //      printf("***site[ %4d ].u = %f\n", n, VLen( siteAngMom[ n ] ) );
      /*       velSum += VLenSq( siteAngMom[ n ] ) + VLenSq( siteVelocity[ n ] ); /\* 3/9/05 this line has been taken into the if block *\/ */
    }    /*     velSum += VLenSq( siteAngMom[ n ] ) + VLenSq( siteVelocity[ n ] ); */
  }
  fprintf( fPtr, "done\n" ); fflush(fPtr);
}

static void SetSiteCoordinates( const int nSite, const real xShift, const real yShift, const real z, const VecRTwo c )
{
  VCopyTwo( siteCoords[ nSite ], c ); 
  siteCoords[ nSite ].x += xShift; // shift x coordinate
  siteCoords[ nSite ].y += yShift; // shift y coordinate
  siteCoords[ nSite ].z = z; // set z coordinate 
}

static void SetSiteFeatures( const int nSite, 
			     const int lipC, 
			     const LipidUnit lipU, 
			     const LipidSpecies lipS, 
			     const LipidUnit lipSkip1,
			     const LipidUnit lipSkip2,
			     const LipidUnit lipSkip3 )
{
  site[ nSite ].inLipid = lipC;
  site[ nSite ].lipidUnit = lipU;
  site[ nSite ].lipidSpecies = lipS;
  site[ nSite ].skipNonbondedInteraction1 = lipSkip1;
  site[ nSite ].skipNonbondedInteraction2 = lipSkip2;
  site[ nSite ].skipNonbondedInteraction3 = lipSkip3;
}

static void SetLipidPair( int *n, // site number
			  int *lipCount, // lipid counter = lipid identifier 
			  VecRTwo c, // lipid xy coordinates 
			  VecRTwo gap, // displacement to avoid overlap 
			  const int lipidSpecies,
			  const VecR region,
			  const real zWaterLayerThickness ) 
{
  real headgroupDistanceFromCenter = INITIAL_BILAYER_THICKNESS / 2; /* initial distance between lipid headgroup and centre of simulation region along z */ 
  real interSiteDistance = 0.2; // [nm] empirical parameter used to space new DOPC model's sites apart
  ++(*lipCount); /* set lipid sites in the "right" monolayer */
  SetSiteCoordinates( *n, 0, 0, headgroupDistanceFromCenter, c );
  SetSiteFeatures( *n, *lipCount, CHOLINE_OR_AMINE, lipidSpecies, PHOSPHATE, -1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, PHOSPHATE, lipidSpecies, CHOLINE_OR_AMINE, GLYCEROL, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, GLYCEROL, lipidSpecies, PHOSPHATE, ESTER_A, ESTER_B );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, ESTER_A, lipidSpecies, GLYCEROL, TAIL_A1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A1, lipidSpecies, ESTER_A, TAIL_A2, -1 );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A2, lipidSpecies, TAIL_A1, TAIL_A3, -1 );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A3, lipidSpecies, TAIL_A2, TAIL_A4, -1 );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A4, lipidSpecies, TAIL_A3, TAIL_A5, -1 );
  ++(*n);
  SetSiteCoordinates( *n, -gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A5, lipidSpecies, TAIL_A4, -1, -1 );
  ++(*n);
  SetSiteCoordinates( *n,  gap.x/4, 0, siteCoords[ *n - 6 ].z, c );
  SetSiteFeatures( *n, *lipCount, ESTER_B, lipidSpecies, GLYCEROL, TAIL_B1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B1, lipidSpecies, ESTER_B, TAIL_B2, -1 );
  ++(*n);
  SetSiteCoordinates( *n, gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B2, lipidSpecies, TAIL_B1, TAIL_B3, -1 );
  ++(*n);
  SetSiteCoordinates( *n, gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B3, lipidSpecies, TAIL_B2, TAIL_B4, -1 );
  ++(*n);
  SetSiteCoordinates( *n, gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B4, lipidSpecies, TAIL_B3, TAIL_B5, -1 );
  ++(*n);
  SetSiteCoordinates( *n, gap.x/4, 0, siteCoords[ *n - 1 ].z - interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B5, lipidSpecies, TAIL_B4, -1, -1 );
  ++(*n);
  
  ++(*lipCount);   /* set lipid sites in the "left" monolayer */
  SetSiteCoordinates( *n, 0, 0, -headgroupDistanceFromCenter, c );
  SetSiteFeatures( *n, *lipCount, CHOLINE_OR_AMINE, lipidSpecies, PHOSPHATE, -1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, 0, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, PHOSPHATE, lipidSpecies, CHOLINE_OR_AMINE, GLYCEROL, -1 );
  ++(*n);
  SetSiteCoordinates( *n,    0,        0, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, GLYCEROL, lipidSpecies, PHOSPHATE, ESTER_A, ESTER_B );
  ++(*n);
  SetSiteCoordinates( *n,    0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, ESTER_A, lipidSpecies, GLYCEROL, TAIL_A1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A1, lipidSpecies, ESTER_A, TAIL_A2, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A2, lipidSpecies, TAIL_A1, TAIL_A3, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A3, lipidSpecies, TAIL_A2, TAIL_A4, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A4, lipidSpecies, TAIL_A3, TAIL_A5, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0, -gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_A5, lipidSpecies, TAIL_A4, -1, -1 );
  ++(*n);
  SetSiteCoordinates( *n,        0,  gap.y/4, siteCoords[ *n - 6 ].z,         c );
  SetSiteFeatures( *n, *lipCount, ESTER_B, lipidSpecies, GLYCEROL, TAIL_B1, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0,  gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B1, lipidSpecies, ESTER_B, TAIL_B2, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0,  gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B2, lipidSpecies, TAIL_B1, TAIL_B3, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0,  gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B3, lipidSpecies, TAIL_B2, TAIL_B4, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0,  gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B4, lipidSpecies, TAIL_B3, TAIL_B5, -1 );
  ++(*n);
  SetSiteCoordinates( *n, 0,  gap.y/4, siteCoords[ *n - 1 ].z + interSiteDistance, c );
  SetSiteFeatures( *n, *lipCount, TAIL_B5, lipidSpecies, TAIL_B4, -1, -1 );
  ++(*n);
}

/*********************************************************************************************************************
 * InitLipidBilayer -- generates from scratch a (non-hydrated) bilayer structure.                                    *
 *                     Currently only single-type bilayers can be generated - to be generalised.                     *
 *********************************************************************************************************************/
void InitLipidBilayer( const int nLipids, const int nDOPCsDSPCs, const int nDOPEs, const VecR region, real *zWaterLayerThickness ) 
{ 
  VecRTwo c, gap, xyRegion;
  int lipCount = 0, n = 0, nx, ny, nLipidsPerEdgeXY;

  // Work out the z-dimension of the water regions above and below the bilayer: 
  *zWaterLayerThickness = 0.5 * ( region.z - INITIAL_BILAYER_THICKNESS );
  *zWaterLayerThickness *= 0.9; /* compressing the thickness to avoid overlaps with lipid headgroups (for a more robust system generation) */

  nLipidsPerEdgeXY = sqrt( nLipids / 2 );
  CheckNLipids( nLipidsPerEdgeXY, nLipids );
  VSetTwo( xyRegion, region.x, region.y ); //  printf("*** xyRegion.x = %f, xyRegion.y = %f\n", xyRegion.x * LENGTH_nm, xyRegion.y * LENGTH_nm);
  VSCopyTwo( gap, 1. / nLipidsPerEdgeXY, xyRegion ); // gap is the lateral edge of each unit cell that will contain a lipid
  for ( ny = 0; ny < nLipidsPerEdgeXY; ny++ ) {
    for ( nx = 0; nx < nLipidsPerEdgeXY; nx++ ) {
      VSetTwo( c, nx + 0.5, ny + 0.5 );
      VMulTwo( c, c, gap );
      VVSAddTwo( c, -0.5, xyRegion );
      if ( nLipids == nDOPCsDSPCs )
	SetLipidPair( &n, &lipCount, c, gap, DOPC, region, *zWaterLayerThickness ); // set lipid pair, one in the top layer and one (the symmetrical) in the bottom layer
      else if ( nLipids == nDOPEs )
	SetLipidPair( &n, &lipCount, c, gap, DOPE, region, *zWaterLayerThickness ); // set lipid pair, one in the top layer and one (the symmetrical) in the bottom layer
      else printf("Error - nLipids must be equal to either nDOPCsDSPCs or nDOPEs (single-component bilayers). For mixtures, InitLipidBilayer must be modified.\n");
    }
  }
}

/*********************************************************************************************************************
 * HydrateLipidBilayer -- adds water to hydrate the bilayer - adapted from [Rapaport, pp. 68-69]                     *
 *                        As usual, it is assumed that the bilayer plane is parallel to the plane defined by         * 
 *                        the x and y dimensions of the simulation region. The water layer thickness is used         *
 *                        to generate the initial coordinates of the membrane system.                                *   
 *********************************************************************************************************************/
void HydrateLipidBilayer( const int nDOPCsDSPCs, const int nDOPEs, const VecR region, const real zWaterLayerThickness, const int nWaters ) 
{
  VecR c, gap, halfWatRegion;
  int n, nx, ny, nz;
  CheckWatInputParameters( nWaters ); // consistency check

  VSet( halfWatRegion, region.x, region.y, zWaterLayerThickness); /* dimensions of the two water regions outside the bilayer */
  VDiv( gap, halfWatRegion, initHalfCellWat );
  n = nDOPCsDSPCs * dopcStruct.nBeads + nDOPEs * dopcStruct.nBeads; 

  // site index corresponding to first element after lipids, i.e., to the first water molecule 
  // first hydrate top bilayer by placing half of the waters in the upper region
  for ( nz = 0; nz < initHalfCellWat.z; nz++ ) {
    for ( ny = 0; ny < initHalfCellWat.y; ny++ ) {
      for ( nx = 0; nx < initHalfCellWat.x; nx++ ) {
        VSet( c, nx + .5, ny + .5, nz + .5 );
        VMul( c, c, gap );
        VVSAdd( c, -0.5, halfWatRegion );
	c.z += 0.5 * ( region.z - zWaterLayerThickness);  // shift upward 
	siteCoords[ n ] = c; 
	site[ n ].inLipid = 0;
	site[ n ].lipidUnit = WATER;
	site[ n ].lipidSpecies = NON_LIPID;
	++n;
      }
    }
  }  
  // now hydrate bottom bilayer by placing half of the waters in the lower region
  for ( nz = 0; nz < initHalfCellWat.z; nz++ ) {
    for ( ny = 0; ny < initHalfCellWat.y; ny++ ) {
      for ( nx = 0; nx < initHalfCellWat.x; nx++ ) {
        VSet( c, nx + .5, ny + .5, nz + .5 );
        VMul( c, c, gap );
        VVSAdd( c, -0.5, halfWatRegion );
	c.z -= 0.5 * ( region.z - zWaterLayerThickness);  // shift downward
	siteCoords[ n ] = c; 
	site[ n ].inLipid = 0;
	site[ n ].lipidUnit = WATER;
	site[ n ].lipidSpecies = NON_LIPID;
	++n;
      }
    }
  }  

  if ( nWaters == 2 ) {
    n = nDOPCsDSPCs * dopcStruct.nBeads + nDOPEs * dopcStruct.nBeads; 
    VSet( siteCoords[ n ], 0,0,0.5);
    VSet( siteCoords[ n+1 ], 0,0, -0.5);
  }
}

